package tt.langchain4j.service;

import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.message.*;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.model.chat.Capability;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.output.TokenUsage;
import dev.langchain4j.rag.AugmentationRequest;
import dev.langchain4j.rag.AugmentationResult;
import dev.langchain4j.rag.query.Metadata;
import dev.langchain4j.service.*;
import dev.langchain4j.service.tool.ToolExecution;
import dev.langchain4j.service.tool.ToolExecutor;
import dev.langchain4j.service.tool.ToolProviderRequest;
import dev.langchain4j.service.tool.ToolProviderResult;

import java.util.*;
import java.util.function.Function;

import static dev.langchain4j.internal.Exceptions.runtime;

public class AiChatServiceFactory extends AiServices<AiChatService> {

    private static final int MAX_SEQUENTIAL_TOOL_EXECUTIONS = 100;


    protected AiChatServiceFactory(AiServiceContext context) {
        super(context);
    }


    public static AiChatServiceFactory builder() {
        return new AiChatServiceFactory(new AiServiceContext(AiChatService.class));
    }


    @Override
    public AiChatService build() {
        return new AiChatServiceImpl();
    }


    private class RequestContext {

        private ServiceRequest request;

        private Object memoryId;

        private UserMessage userMessage;

        private AugmentationResult augmentationResult;

        private List<ToolSpecification> toolSpecifications;

        private Map<String, ToolExecutor> toolExecutors;

        private List<ChatMessage> messages;

        public RequestContext(ServiceRequest request) {
            this.request = request;
        }

    }


    private class AiChatServiceImpl implements AiChatService {

        @Override
        public TokenStream chatWithStream(ServiceRequest request) {
            final RequestContext requestContext = new RequestContext(request);
            beforeChat(requestContext);
            System.out.println("onChat");
            return new AiServiceTokenStream(
                    requestContext.messages,
                    requestContext.toolSpecifications,
                    requestContext.toolExecutors,
                    requestContext.augmentationResult != null ? requestContext.augmentationResult.contents() : null,
                    context,
                    requestContext.memoryId
            );
        }

        @Override
        public Result<AiMessage> chat(ServiceRequest request) {
            final RequestContext requestContext = new RequestContext(request);
            beforeChat(requestContext);
            ChatRequestParameters parameters = ChatRequestParameters.builder()
                    .toolSpecifications(requestContext.toolSpecifications)
                    .build();
            ChatRequest chatRequest = ChatRequest.builder()
                    .messages(requestContext.messages)
                    .parameters(parameters)
                    .build();
            ChatResponse chatResponse = context.chatModel.chat(chatRequest);
            TokenUsage tokenUsageAccumulator = chatResponse.metadata().tokenUsage();
            //
            int executionsLeft = MAX_SEQUENTIAL_TOOL_EXECUTIONS;
            List<ToolExecution> toolExecutions = new ArrayList<>();
            while (true) {
                if (executionsLeft-- == 0) {
                    throw runtime("Something is wrong, exceeded %s sequential tool executions",
                            MAX_SEQUENTIAL_TOOL_EXECUTIONS);
                }

                AiMessage aiMessage = chatResponse.aiMessage();

                if (context.hasChatMemory()) {
                    context.chatMemory(requestContext.memoryId).add(aiMessage);
                } else {
                    requestContext.messages = new ArrayList<>(requestContext.messages);
                    requestContext.messages.add(aiMessage);
                }

                if (!aiMessage.hasToolExecutionRequests()) {
                    break;
                }

                for (ToolExecutionRequest toolExecutionRequest : aiMessage.toolExecutionRequests()) {
                    ToolExecutor toolExecutor = requestContext.toolExecutors.get(toolExecutionRequest.name());
                    String toolExecutionResult = toolExecutor.execute(toolExecutionRequest, requestContext.memoryId);
                    toolExecutions.add(ToolExecution.builder()
                            .request(toolExecutionRequest)
                            .result(toolExecutionResult)
                            .build());
                    ToolExecutionResultMessage toolExecutionResultMessage = ToolExecutionResultMessage.from(
                            toolExecutionRequest,
                            toolExecutionResult
                    );
                    if (context.hasChatMemory()) {
                        context.chatMemory(requestContext.memoryId).add(toolExecutionResultMessage);
                    } else {
                        requestContext.messages.add(toolExecutionResultMessage);
                    }
                }

                if (context.hasChatMemory()) {
                    requestContext.messages = context.chatMemory(requestContext.memoryId).messages();
                }

                chatRequest = ChatRequest.builder()
                        .messages(requestContext.messages)
                        .parameters(parameters)
                        .build();

                System.out.println("onChat");
                chatResponse = context.chatModel.chat(chatRequest);

                tokenUsageAccumulator = TokenUsage.sum(tokenUsageAccumulator, chatResponse.metadata().tokenUsage());
            }

            //返回
            return Result.<AiMessage>builder()
                    .content(chatResponse.aiMessage())
                    .tokenUsage(tokenUsageAccumulator)
                    .sources(requestContext.augmentationResult == null ? null : requestContext.augmentationResult.contents())
                    .finishReason(chatResponse.metadata().finishReason())
                    .toolExecutions(toolExecutions)
                    .build();
        }



        /**
         * 复制
         * @see //dev.langchain4j.service.DefaultAiServices
         */
        private boolean supportsJsonSchema() {
            return context.chatModel != null
                    && context.chatModel.supportedCapabilities().contains(Capability.RESPONSE_FORMAT_JSON_SCHEMA);
        }

    }

    private void beforeChat( RequestContext requestContext ){
        requestContext.memoryId = Optional.ofNullable(requestContext.request.getMemoryId()).orElse(DEFAULT);
        final Optional<SystemMessage> systemMessage = Optional.ofNullable(requestContext.request.getSystemMessage());
        //first
        requestContext.userMessage = requestContext.request.getUserMessage();
        //rag
        if (context.retrievalAugmentor != null) {
            List<ChatMessage> chatMemory = context.hasChatMemory()
                    ? context.chatMemory(requestContext.memoryId).messages()
                    : null;
            Metadata metadata = Metadata.from(requestContext.userMessage, requestContext.memoryId, chatMemory);
            AugmentationRequest augmentationRequest = new AugmentationRequest(requestContext.userMessage, metadata);
            requestContext.augmentationResult = context.retrievalAugmentor.augment(augmentationRequest);
            //替换
            requestContext.userMessage = (UserMessage) requestContext.augmentationResult.chatMessage();
        }
        //拼接媒体内容
        final List<Content> medias = requestContext.request.getMedias();
        if ( null != medias && medias.size() > 0 ) {
            final List<Content> contents = new ArrayList<>(requestContext.userMessage.contents());
            contents.addAll(medias);
            if ( null != requestContext.userMessage.name() ) {
                requestContext.userMessage = UserMessage.userMessage(
                        requestContext.userMessage.name() ,
                        contents
                );
            } else {
                requestContext.userMessage = UserMessage.userMessage(
                        contents
                );
            }
        }


        //存入
        if (context.hasChatMemory()) {
            ChatMemory chatMemory = context.chatMemory(requestContext.memoryId);
            systemMessage.ifPresent(chatMemory::add);
            chatMemory.add(requestContext.userMessage);
        }

        if (context.hasChatMemory()) {
            requestContext.messages = context.chatMemory(requestContext.memoryId).messages();
        } else {
            requestContext.messages = new ArrayList<>();
            systemMessage.ifPresent(requestContext.messages::add);
            requestContext.messages.add(requestContext.userMessage);
        }

        //tools
        requestContext.toolSpecifications = context.toolSpecifications;
        requestContext.toolExecutors = context.toolExecutors;
        if (context.toolProvider != null) {
            requestContext.toolSpecifications = new ArrayList<>();
            requestContext.toolExecutors = new HashMap<>();
            ToolProviderRequest toolProviderRequest = new ToolProviderRequest(requestContext.memoryId,requestContext.userMessage);
            ToolProviderResult toolProviderResult = context.toolProvider.provideTools(toolProviderRequest);
            if (toolProviderResult != null) {
                Map<ToolSpecification, ToolExecutor> tools = toolProviderResult.tools();
                for (ToolSpecification toolSpecification : tools.keySet()) {
                    requestContext.toolSpecifications.add(toolSpecification);
                    requestContext.toolExecutors.put(toolSpecification.name(), tools.get(toolSpecification));
                }
            }
        }

    }



    //取消部分功能



    @Override
    @Deprecated
    public AiServices<AiChatService> systemMessageProvider(Function<Object, String> systemMessageProvider) {
        return this;
    }



}
